package com.sui.bigdata.sml.web.misc.k8s;

import com.feidee.fdcommon.configuration.CustomConfiguration;
import com.sui.bigdata.sml.web.configuration.KubeConfig;
import com.sui.bigdata.sml.web.util.Constants;
import com.sui.bigdata.sml.web.util.ModelType;
import io.kubernetes.client.custom.IntOrString;
import io.kubernetes.client.models.*;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * @author songhaicheng
 * @date 2020/3/11 14:35
 * @description
 * @reviewer
 */
public class KubeComponentMaker {
    private static Map<String, String> nodeConfig = new HashMap<>();
    private static final String TF_NAME = "c0";
    private static final String SML_NAME = "c1";
    private static final V1ContainerPort TF_PORT = new V1ContainerPort().containerPort(8500);
    private static final V1ContainerPort SML_PORT = new V1ContainerPort().containerPort(8000);
    private static final V1ServicePort SERVICE_PORT = new V1ServicePortBuilder()
            .withName("http")
            .withPort(80)
            .withTargetPort(new IntOrString(8000))
            .withProtocol("TCP")
            .build();

    static {
        nodeConfig.put("nodetype", KubeConfig.NAMESPACE);
    }

    private String name;
    private Map<String, String> labelCofig;
    private Map<String, String> ingressConfig;
    private ServingInfo servingInfo;
    private V1Deployment deployment;
    private V1Service service;
    private ExtensionsV1beta1Ingress ingress;

    public KubeComponentMaker(ServingInfo servingInfo) {
        this.servingInfo = servingInfo;
        name = servingInfo.getServingToken();
        labelCofig = new HashMap<>();
        labelCofig.put("app", KubeConfig.NAME_PREFIX + name);
        ingressConfig = new HashMap<>();
        ingressConfig.put("nginx.ingress.kubernetes.io/rewrite-target", "/model/" + servingInfo.getServingToken());
    }

    private void makeDeployment() {
        List<ServingModel> tfModel = servingInfo.getModels().stream()
                .filter(m -> ModelType.getType(m.getModelType()) == ModelType.TF)
                .collect(Collectors.toList());
        List<V1Container> containers = new ArrayList<>();
        V1Volume volume = null;
        if (tfModel.size() == 1) {
            //tensorflow模型配置参数
            String tfArgs = CustomConfiguration.getString(makeTensorflowArgsKey(name),
                    "--tensorflow_intra_op_parallelism=10," +
                            "--tensorflow_inter_op_parallelism=10," +
                            "--enable_batching=true," +
                            "--flush_filesystem_caches=true");

            volume = new V1VolumeBuilder()
                    .withName("tf-model-path")
                    .withHostPath(new V1HostPathVolumeSource().path(tfModel.get(0).getHdfsPath()))
                    .build();
            V1VolumeMount tfModelMount = new V1VolumeMountBuilder()
                    .withName(volume.getName())
                    .withMountPath("/models/model/1")
                    .build();
            V1Container tfContainer = new V1ContainerBuilder()
                    .withName(TF_NAME)
                    .withImage(KubeConfig.TF_IMAGE)
                    .withImagePullPolicy("IfNotPresent")
                    .withPorts(TF_PORT)
                    .withVolumeMounts(tfModelMount)
                    .withArgs(tfArgs.split(","))
                    .build();
            containers.add(tfContainer);
        }
        // 服务是否可用探针，由统一定义的启动超时时间 T 进行分配，延迟时间 = T * 0.4，最大允许尝试 3 次，每次间隔时间 = T * 0.2
        int startTimeout = CustomConfiguration.getInt(Constants.KUBE_SML_READINESS_SECONDS_TIMEOUT, 60);
        V1Probe readinessProbe = new V1ProbeBuilder()
                .withHttpGet(new V1HTTPGetActionBuilder()
                        .withPath("/model/" + servingInfo.getServingToken())
                        .withPort(new IntOrString(SML_PORT.getContainerPort()))
                        .build())
                .withInitialDelaySeconds(startTimeout / 5 * 2)
                .withPeriodSeconds(startTimeout / 5)
                .withFailureThreshold(3)
                .build();
        // 服务是否可用探针，用于不可用时重启容器
        V1Probe livenessProbe = new V1ProbeBuilder()
                .withHttpGet(new V1HTTPGetActionBuilder()
                        .withPath("/model/" + servingInfo.getServingToken())
                        .withPort(new IntOrString(SML_PORT.getContainerPort()))
                        .build())
                .withInitialDelaySeconds(startTimeout / 5 * 2)
                .withPeriodSeconds(CustomConfiguration.getInt(Constants.KUBE_SML_LIVENESS_PROBE_SECONDS_PERIOD, 60))
                .withFailureThreshold(CustomConfiguration.getInt(Constants.KUBE_SML_LIVENESS_PROBE_FAILURE_THRESHOLD, 3))
                .build();
        V1Container smlContainer = new V1ContainerBuilder()
                .withName(SML_NAME)
                .withImage(KubeConfig.SML_IMAGE)
                .withImagePullPolicy("Always")
                .withCommand("java", "-jar",CustomConfiguration.getString(makeJvmArgsKey(name), "-Xmixed"), "sml-model.jar")
                .withArgs(servingInfo.toString())
                .withPorts(SML_PORT)
                .withReadinessProbe(readinessProbe)
                .withLivenessProbe(livenessProbe)
                .build();
        containers.add(smlContainer);
        V1Volume[] volumes = volume == null ? new V1Volume[0] : new V1Volume[]{volume};
        this.deployment = new V1DeploymentBuilder()
                .withApiVersion("apps/v1")
                .withKind("Deployment")
                .withNewMetadata()
                .withName(KubeConfig.NAME_PREFIX + name)
                .withNamespace(KubeConfig.NAMESPACE)
                .endMetadata()
                .withNewSpec()
                .withReplicas(CustomConfiguration.getInt(makeDeploymentReplicasKey(name), 2))
                .withNewSelector()
                .withMatchLabels(labelCofig)
                .endSelector()
                .withNewTemplate()
                .withNewMetadata()
                .withLabels(labelCofig)
                .endMetadata()
                .withNewSpec()
                .withNodeSelector(nodeConfig)
                .addToContainers(containers.toArray(new V1Container[0]))
                .withVolumes(volumes)
                .withImagePullSecrets(new V1LocalObjectReference().name("registry-secret"))
                .endSpec()
                .endTemplate()
                .endSpec()
                .build();
    }

    private void makeService() {
        this.service =  new V1ServiceBuilder()
                .withApiVersion("v1")
                .withKind("Service")
                .withNewMetadata()
                .withName("sml-model-" + servingInfo.getServingToken())
                .withNamespace(KubeConfig.NAMESPACE)
                .endMetadata()
                .withNewSpec()
                .withSelector(labelCofig)
                .withPorts(SERVICE_PORT)
                .endSpec()
                .build();
    }

    private void makeIngress() {
        ExtensionsV1beta1IngressBackend backend = new ExtensionsV1beta1IngressBackendBuilder()
                .withServiceName(KubeConfig.NAME_PREFIX + name)
                .withServicePort(new IntOrString(80))
                .build();
        ExtensionsV1beta1HTTPIngressPath path = new ExtensionsV1beta1HTTPIngressPathBuilder()
                .withPath(ingressConfig.get("nginx.ingress.kubernetes.io/rewrite-target") + "(/|$)(.*)")
                .withBackend(backend)
                .build();
        ExtensionsV1beta1HTTPIngressRuleValue http = new ExtensionsV1beta1HTTPIngressRuleValueBuilder()
                .withPaths(path)
                .build();
        ExtensionsV1beta1IngressRule rule = new ExtensionsV1beta1IngressRuleBuilder()
                .withHttp(http)
                .build();

        this.ingress = new ExtensionsV1beta1IngressBuilder()
                .withApiVersion("extensions/v1beta1")
                .withKind("Ingress")
                .withNewMetadata()
                .withName(KubeConfig.INGRESS_NAME_PREFIX + name)
                .withNamespace(KubeConfig.NAMESPACE)
                .withAnnotations(ingressConfig)
                .endMetadata()
                .withNewSpec()
                .withRules(rule)
                .endSpec()
                .build();
    }

    public V1Deployment getDeployment() {
        if (deployment == null) {
            makeDeployment();
        }
        return deployment;
    }

    public V1Service getService() {
        if (service == null) {
            makeService();
        }
        return service;
    }

    public ExtensionsV1beta1Ingress getIngress() {
        if (ingress == null) {
            makeIngress();
        }
        return ingress;
    }

    public String makeJvmArgsKey(String servingToken) {
        return servingToken + ".jvm.args";
    }

    public String makeDeploymentReplicasKey(String servingToken) {
        return servingToken + ".deployment.replicas";
    }

    public String makeTensorflowArgsKey(String servingToken) {
        return servingToken + ".tensorflow.args";
    }
}
